1256E - Yet Another Division Into Teams - CodeForces Solution

#include "bits/stdc++.h"
using namespace std;
typedef long long ll;
#define INF 1e12

const int N = 2e5 + 7;
int n, idx[N];
pair<int, int> a[N];
ll dp[N];

ll rec(int i)
{
    if (i == n)
        return 0;
    if(n - i < 3)
        return INF;
    ll &ret = dp[i];
    if (~ret)
        return ret;
    ret = a[n - 1].first - a[i].first;
    int low = i + 3, high = i + 5;
    for(low; low <= high; ++low)
    {
        ll ans = rec(low), cur = (a[low - 1].first - a[i].first);
        ret = min(ret, ans + cur);
    }
    return ret;
}

void solve()
{
    cin >> n;
    for (int i = 0; i < n; ++i)
        cin >> a[i].first, a[i].second = i;
    sort(a, a + n);
    memset(dp, -1, sizeof(dp));
    ll ans = rec(0);
    cout << ans << ' ';
    int cnt = 1, lst = 0;
    idx[a[0].second] = cnt;
    // build
    int number = 1;
    for (int i = 1; i < n; ++i)
    {
        if (number > 2 and ans - (a[i - 1].first - a[lst].first) == dp[i])
        {
            cnt++;
            ans -= (a[i - 1].first - a[lst].first);
            lst = i;
            number = 1;
        }
        else
            number++;
        idx[a[i].second] = cnt;
    }
    cout << cnt << '\n';
    for (int i = 0; i < n; ++i)
        cout << idx[i] << ' ';
    cout << '\n';
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int tc = 1;
    // cin >> tc;
    for (int i = 1; i <= tc; ++i)
        solve();
}